import re
import pandas as pd
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
from itertools import combinations
from tqdm import tqdm
import concurrent.futures

# --- 您可以调整这里的参数 ---
NUM_CANDIDATES_TO_CONSIDER = 15
PORTFOLIO_SIZE = 8
MIN_TRADES = 50

# --- Matplotlib 设置 ---
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False


def parse_optimization_results(file_path):
    """从日志文件中解析所有的OptimizationResult行。"""
    # (此函数无需改动，代码省略以保持简洁)
    results = []
    try:
        with open(file_path, 'r', encoding='UTF-8') as f:
            content = f.read()
    except FileNotFoundError:
        print(f"错误: 文件 '{file_path}' 未找到。请确保文件路径正确。")
        return None

    pattern = re.compile(
        r"OptimizationResult\(symbol=(\w+), "
        r"finalBalance=([-\d.]+), "
        r"totalReturn=([-\d.]+), "
        r"maxDrawdown=([-\d.]+), "
        r"profitFactor=([-\d.]+), "
        r"totalTrades=(\d+), "
        r"winCount=(\d+), "
        r"lossCount=(\d+), "
        r"winRate=([-\d.]+), "
        r".*?, " # Skip startDate, endDate
        r"dailyReturns=\[(.*?)\]" # Non-greedy match for the list content
    )

    matches = pattern.finditer(content)

    for match in matches:
        try:
            returns_str = match.group(10)
            if returns_str:
                daily_returns = [float(r) for r in returns_str.split(', ') if r]
            else:
                daily_returns = []

            results.append({
                'symbol': match.group(1),
                'totalReturn': float(match.group(3)),
                'maxDrawdown': float(match.group(4)),
                'profitFactor': float(match.group(5)),
                'totalTrades': int(match.group(6)),
                'winRate': float(match.group(9)),
                'dailyReturns': daily_returns
            })
        except (ValueError, IndexError) as e:
            symbol_name = match.group(1) if match else "未知"
            print(f"解析币种 '{symbol_name}' 的数据时出错: {e}. 跳过该行。")

    if not results:
        print("错误：未能在文件中找到任何有效的'OptimizationResult'数据。请检查日志文件格式和正则表达式。")
        return None

    return pd.DataFrame(results)


def calculate_portfolio_performance(df, portfolio_symbols):
    """计算指定币种组合的表现。"""
    # (此函数无需改动，代码省略以保持简洁)
    portfolio_returns_df = pd.DataFrame()
    for symbol in portfolio_symbols:
        returns_series_list = df.loc[df['symbol'] == symbol, 'dailyReturns'].values
        if len(returns_series_list) > 0:
            portfolio_returns_df[symbol] = pd.Series(returns_series_list[0])

    if portfolio_returns_df.empty:
        return 0, 1, pd.Series()

    portfolio_returns_df.fillna(0, inplace=True)
    portfolio_daily_returns = portfolio_returns_df.mean(axis=1)

    initial_capital = 1.0
    equity_curve = (1 + portfolio_daily_returns).cumprod() * initial_capital

    if equity_curve.empty:
        return 0, 1, pd.Series()

    running_max = equity_curve.cummax()
    drawdown = (equity_curve - running_max) / running_max
    max_drawdown = abs(drawdown.min()) if not drawdown.empty else 0
    total_return = equity_curve.iloc[-1] - 1 if not equity_curve.empty else 0

    return total_return, max_drawdown, equity_curve


def evaluate_combination_worker(combo, df):
    """
    这是进程池中每个子进程要执行的任务。
    【核心改动 1】: 增加计算 'returnToDrawdownRatio'
    """
    total_return, max_drawdown, _ = calculate_portfolio_performance(df, combo)

    # 计算收益回撤比，处理 max_drawdown 为0的特殊情况
    # 加上一个极小值避免除以零
    ratio = total_return / (max_drawdown + 1e-9)

    return {
        'portfolio': ", ".join(combo),
        'portfolio_totalReturn': total_return,
        'portfolio_maxDrawdown': max_drawdown,
        'returnToDrawdownRatio': ratio # 新增的指标
    }


def find_best_portfolio_balanced(df): # 函数名改为 balanced
    """
    【多进程版本】遍历组合，找到收益和回撤最平衡的组合。
    """
    # 1. 过滤和质量评分 (这部分不变, 仍然用于筛选出综合素质不错的候选币)
    df_filtered = df[(df['totalTrades'] >= MIN_TRADES) & (df['profitFactor'] >= 1.0)].copy()
    df_filtered['qualityScore'] = (
            df_filtered['totalReturn'].rank(pct=True, ascending=False) +
            df_filtered['maxDrawdown'].rank(pct=True, ascending=True) +
            df_filtered['profitFactor'].rank(pct=True, ascending=False) +
            df_filtered['winRate'].rank(pct=True, ascending=False)
    )
    candidates = df_filtered.sort_values('qualityScore', ascending=False).head(NUM_CANDIDATES_TO_CONSIDER)
    candidate_symbols = ['SUIUSDT', 'ETHUSDT', 'DOTUSDT', 'INJUSDT', 'DOGEUSDT', 'GRTUSDT',
                         'LINKUSDT', 'ADAUSDT', 'ALGOUSDT', 'OPUSDT', 'SOLUSDT', 'AVAXUSDT',
                         'ARBUSDT', 'TIAUSDT', 'ONDOUSDT', ]

    if len(candidate_symbols) < PORTFOLIO_SIZE:
        print(f"警告：高质量候选币种数量({len(candidate_symbols)})少于期望的组合数量({PORTFOLIO_SIZE})。")
        return pd.DataFrame(), candidates

    print(f"\n从以下 {len(candidate_symbols)} 个高质量候选币种中进行组合筛选:")
    print(", ".join(candidate_symbols))

    portfolio_combinations = list(combinations(candidate_symbols, PORTFOLIO_SIZE))
    print(f"\n总共需要评估 {len(portfolio_combinations)} 种组合...")

    # 3. 使用多进程并行计算
    portfolio_results = []
    with concurrent.futures.ProcessPoolExecutor() as executor:
        futures = [executor.submit(evaluate_combination_worker, combo, df) for combo in portfolio_combinations]
        for future in tqdm(concurrent.futures.as_completed(futures), total=len(futures), desc="正在评估投资组合"):
            try:
                result = future.result()
                portfolio_results.append(result)
            except Exception as e:
                print(f"一个子任务计算出错: {e}")

    if not portfolio_results:
        return pd.DataFrame(), candidates

    results_df = pd.DataFrame(portfolio_results)

    # 【核心改动 2】: 更改排序逻辑，按 'returnToDrawdownRatio' 降序排列
    results_df = results_df.sort_values('returnToDrawdownRatio', ascending=False)

    return results_df, candidates


def create_correlation_matrix(df, symbols):
    """根据每日收益率构建相关性矩阵。"""
    # (此函数无需改动，代码省略以保持简洁)
    returns_dict = {
        symbol: df.loc[df['symbol'] == symbol, 'dailyReturns'].iloc[0]
        for symbol in symbols if symbol in df['symbol'].values
    }
    returns_df = pd.DataFrame(dict([(k, pd.Series(v)) for k, v in returns_dict.items()]))
    returns_df.fillna(0, inplace=True)
    return returns_df.corr()


# --- 主程序 ---
if __name__ == "__main__":
    file_path = '回测结果 (2).txt'
    raw_df = parse_optimization_results(file_path)

    if raw_df is not None and not raw_df.empty:
        # 调用新的、平衡的函数
        best_portfolios_df, candidates_df = find_best_portfolio_balanced(raw_df)

        if not best_portfolios_df.empty:
            # 【核心改动 3】: 更新打印输出，显示新的排序指标
            print("\n==================== 综合表现最佳的Top 5投资组合 (按收益回撤比排名) ====================")
            display_df = best_portfolios_df.head(5).copy()
            display_df['portfolio_totalReturn'] = display_df['portfolio_totalReturn'].apply(lambda x: f"{x:.2%}")
            display_df['portfolio_maxDrawdown'] = display_df['portfolio_maxDrawdown'].apply(lambda x: f"{x:.2%}")
            # 格式化收益回撤比，保留两位小数
            display_df['returnToDrawdownRatio'] = display_df['returnToDrawdownRatio'].apply(lambda x: f"{x:.2f}")
            print(display_df.to_string())

            # 后续的深度分析和绘图部分无需改动，它会自动分析新的最佳组合
            best_portfolio_symbols_str = best_portfolios_df.iloc[0]['portfolio']
            best_portfolio_symbols = [s.strip() for s in best_portfolio_symbols_str.split(',')]

            print("\n\n==================== 最佳投资组合深度分析 ====================")
            print(f"币种组合: {best_portfolio_symbols_str}")

            best_return, best_drawdown, best_equity_curve = calculate_portfolio_performance(raw_df, best_portfolio_symbols)
            print(f"组合总回报: {best_return:.2%}")
            print(f"组合最大回撤: {best_drawdown:.2%}")
            # 计算并打印最佳组合的收益回撤比
            best_ratio = best_return / (best_drawdown + 1e-9)
            print(f"收益回撤比: {best_ratio:.2f}")

            # (绘图部分无需改动，代码省略)
            plt.figure(figsize=(15, 8))
            best_equity_curve.plot(label='核心投资组合权益曲线', color='red', linewidth=2.5)

            for symbol in best_portfolio_symbols:
                returns = raw_df[raw_df['symbol'] == symbol]['dailyReturns'].iloc[0]
                equity = (1 + pd.Series(returns)).cumprod()
                equity.plot(label=symbol, alpha=0.4, linestyle='--')

            plt.title('最佳投资组合 vs. 单个币种权益曲线', fontsize=16)
            plt.ylabel('权益价值 (初始为1)')
            plt.xlabel('交易日')
            plt.legend()
            plt.grid(True)
            plt.show()

            corr_matrix = create_correlation_matrix(raw_df, best_portfolio_symbols)
            plt.figure(figsize=(10, 8))
            sns.heatmap(corr_matrix, annot=True, cmap='coolwarm', fmt=".2f", vmin=-1, vmax=1)
            plt.title('最佳投资组合内部相关性', fontsize=16)
            plt.show()

        else:
            print("\n未能计算出任何投资组合，请检查候选币种数量和组合大小设置。")